package com.learn.security.utils.asymmetric_encryption;

import com.learn.security.utils.BaseUtil;

import javax.crypto.Cipher;
import java.math.BigInteger;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.*;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA加密工具类
 * <p>
 * 既能进行签名，也能进行加密通信
 */
public class RSAUtil extends BaseUtil {

    public static final String KEY_ALGORITHM = "RSA";
    public static final String SIGNATURE_ALGORITHM = "MD5withRSA";

    /**
     * 默认密钥字节数
     */
    public static final int KEY_SIZE = 1024;

    public static final String PUBLIC_KEY = "RSAPublicKey";
    public static final String PRIVATE_KEY = "RSAPrivateKey";


    /**
     * 初始化密钥
     *
     * @param keySize 密钥字节数
     * @return 装秘钥和公钥的map
     */
    public static Map<String, Object> initKey(int keySize) {
        Map<String, Object> keyMap = new HashMap<>();

        try {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(KEY_ALGORITHM);
            keyPairGen.initialize(keySize);

            KeyPair keyPair = keyPairGen.generateKeyPair();

            // 公钥
            RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();

            // 私钥
            RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();

            keyMap.put(PUBLIC_KEY, publicKey);
            keyMap.put(PRIVATE_KEY, privateKey);

        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return keyMap;
    }

    /**
     * 初始化密钥
     * 默认密钥字节数为1024
     *
     * @return 装秘钥和公钥的map
     */
    public static Map<String, Object> initKey() {
        return initKey(KEY_SIZE);
    }

    /**
     * 取得私钥--原始值
     *
     * @param keyMap 装秘钥和公钥的map
     * @return 私钥--base64加密的字符
     */
    public static RSAPrivateKey getPrivateKeyOriginal(Map<String, Object> keyMap) {
        return (RSAPrivateKey) keyMap.get(PRIVATE_KEY);
    }

    /**
     * 取得私钥--base64加密
     *
     * @param keyMap 装秘钥和公钥的map
     * @return 私钥--base64加密的字符
     */
    public static String getPrivateKey(Map<String, Object> keyMap) {
        return encryptBASE64(getPrivateKeyOriginal(keyMap).getEncoded());
    }

    /**
     * 取得公钥--原始值
     *
     * @param keyMap 装秘钥和公钥的map
     * @return 公钥--RSAPublicKey
     */
    public static RSAPublicKey getPublicKeyOriginal(Map<String, Object> keyMap) {
        return (RSAPublicKey) keyMap.get(PUBLIC_KEY);
    }

    /**
     * 取得公钥--base64加密
     *
     * @param keyMap 装秘钥和公钥的map
     * @return 公钥--base64加密的字符
     */
    public static String getPublicKey(Map<String, Object> keyMap) {
        return encryptBASE64(getPublicKeyOriginal(keyMap).getEncoded());
    }

    /**
     * 根据 n，e 获取公钥
     *
     * @param modulus        n
     * @param publicExponent e
     * @return 公钥--base64加密的字符
     */
    public static String getPublicKey(byte[] modulus, byte[] publicExponent) {

        RSAPublicKeySpec pubKeySpec = new RSAPublicKeySpec(new BigInteger(modulus), new BigInteger(publicExponent));

        try {
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            RSAPublicKey publicKey = (RSAPublicKey) keyFactory.generatePublic(pubKeySpec);
            return encryptBASE64(publicKey.getEncoded());
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 根据私钥获取Modulus n
     *
     * @param key 私钥
     * @return Modulus
     */
    public static BigInteger getModulus(RSAPrivateKey key) {
        return key.getModulus();
    }

    /**
     * 根据公钥获取Modulus n
     *
     * @param key 公钥
     * @return Modulus
     */
    public static BigInteger getModulus(RSAPublicKey key) {
        return key.getModulus();
    }

    /**
     * 根据私钥获取PrivateExponent d
     *
     * @param key 私钥
     * @return PrivateExponent
     */
    public static BigInteger getPrivateExponent(RSAPrivateKey key) {
        return key.getPrivateExponent();
    }

    /**
     * 根据公钥获取PublicExponent e
     *
     * @param key 私钥
     * @return PublicExponent
     */
    public static BigInteger getPublicExponent(RSAPublicKey key) {
        return key.getPublicExponent();
    }


    /**
     * 加密--用私钥加密
     *
     * @param data 要加密byte[]
     * @param key  私钥 base64加密的字符
     * @return 加密后的byte[]
     */
    public static byte[] encryptByPrivateKey(byte[] data, String key) {
        try {
            // 对密钥解密
            byte[] keyBytes = decryptBASE64(key);

            if (null != keyBytes) {
                // 取得私钥
                PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(keyBytes);
                KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
                Key privateKey = keyFactory.generatePrivate(pkcs8KeySpec);

                // 对数据加密
                Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
                cipher.init(Cipher.ENCRYPT_MODE, privateKey);

                return cipher.doFinal(data);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }


    /**
     * 加密--用公钥加密
     *
     * @param data 要加密byte[]
     * @param key  公钥 base64加密的字符
     * @return 加密后的byte[]
     */
    public static byte[] encryptByPublicKey(byte[] data, String key) {
        // 对公钥解密
        byte[] keyBytes = decryptBASE64(key);

        try {
            if (null != keyBytes) {
                // 取得公钥
                X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
                KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
                Key publicKey = keyFactory.generatePublic(x509KeySpec);

                // 对数据加密
                Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
                cipher.init(Cipher.ENCRYPT_MODE, publicKey);

                return cipher.doFinal(data);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 解密--用私钥解密
     *
     * @param data 要解密的byte[]
     * @param key  私钥 base64加密的字符
     * @return 解密后的byte[]
     */
    public static byte[] decryptByPrivateKey(byte[] data, String key) {
        // 对密钥解密
        byte[] keyBytes = decryptBASE64(key);

        try {
            if (null != keyBytes) {
                // 取得私钥
                PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(keyBytes);
                KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
                Key privateKey = keyFactory.generatePrivate(pkcs8KeySpec);

                // 对数据解密
                Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
                cipher.init(Cipher.DECRYPT_MODE, privateKey);

                return cipher.doFinal(data);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 解密--用公钥解密
     *
     * @param data 要解密的byte[]
     * @param key  私钥 base64加密的字符
     * @return 解密后的byte[]
     */
    public static byte[] decryptByPublicKey(byte[] data, String key) {
        // 对密钥解密
        byte[] keyBytes = decryptBASE64(key);

        try {
            if (null != keyBytes) {
                // 取得公钥
                X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
                KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
                Key publicKey = keyFactory.generatePublic(x509KeySpec);

                // 对数据解密
                Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
                cipher.init(Cipher.DECRYPT_MODE, publicKey);

                return cipher.doFinal(data);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 用私钥对信息生成数字签名
     *
     * @param data       加密数据
     * @param privateKey 私钥  base64加密的字符
     * @return 签名--base64加密的字符
     */
    public static String sign(byte[] data, String privateKey) {
        // 解密由base64编码的私钥
        byte[] keyBytes = decryptBASE64(privateKey);

        try {
            if (null != keyBytes) {
                // 构造PKCS8EncodedKeySpec对象
                PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(keyBytes);

                // KEY_ALGORITHM 指定的加密算法
                KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);

                // 取私钥匙对象
                PrivateKey priKey = keyFactory.generatePrivate(pkcs8KeySpec);

                // 用私钥对信息生成数字签名
                Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
                signature.initSign(priKey);
                signature.update(data);

                return encryptBASE64(signature.sign());
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    /**
     * 校验数字签名
     *
     * @param data      加密数据
     * @param publicKey 公钥
     * @param sign      数字签名
     * @return 校验成功返回true 失败返回false
     */
    public static Boolean verify(byte[] data, String publicKey, String sign) {

        // 解密由base64编码的公钥
        byte[] keyBytes = decryptBASE64(publicKey);

        try {
            if (null != keyBytes) {
                // 构造X509EncodedKeySpec对象
                X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);

                // KEY_ALGORITHM 指定的加密算法
                KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);

                // 取公钥匙对象
                PublicKey pubKey = keyFactory.generatePublic(keySpec);

                Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
                signature.initVerify(pubKey);
                signature.update(data);

                // 验证签名是否正常
                return signature.verify(decryptBASE64(sign));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }
}
